import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt

from CNN import CNN
from read_data import DataLoader


def train_model(model, device, data_directory, num_epochs=10, batch_sz=16, learning_rate=1e-5):
    # 数据加载
    dataset = DataLoader(data_directory)
    train_data_loader = torch.utils.data.DataLoader(dataset=dataset, batch_size=batch_sz, shuffle=True)

    # 优化器设置
    optimizer_instance = optim.RMSprop(model.parameters(), lr=learning_rate, weight_decay=1e-8, momentum=0.9)

    # 损失函数定义
    loss_fn = nn.CrossEntropyLoss()

    # 初始化最佳损失为无穷大
    best_loss_value = float('inf')
    epoch_losses_history = []

    # 训练循环
    for epoch_idx in range(num_epochs):
        model.train()  # 设置模型为训练模式
        print(f"Epoch {epoch_idx + 1}/{num_epochs}")

        cumulative_loss = 0.0
        for images, labels in train_data_loader:
            optimizer_instance.zero_grad()
            images = images.to(device=device, dtype=torch.float32)
            labels = labels.to(device=device, dtype=torch.long)

            predictions = model(images)
            loss = loss_fn(predictions, labels)
            print('Training Loss', loss.item())
            cumulative_loss += loss.item()

            # 更新最佳损失并保存模型
            if loss < best_loss_value:
                best_loss_value = loss
                torch.save(model.state_dict(), 'best_model.pth')

            # 反向传播和优化步骤
            loss.backward()
            optimizer_instance.step()

        # 计算并记录平均损失
        average_loss = cumulative_loss / len(train_data_loader)
        epoch_losses_history.append(average_loss)
        print(
            f"Epoch {epoch_idx + 1}/{num_epochs} completed, Avg Loss: {average_loss:.4f}, Best Loss: {best_loss_value:.4f}")

    # 绘制损失下降曲线
    plt.plot(range(1, num_epochs + 1), epoch_losses_history, label="Training Loss", color='blue')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title('Training Loss Decrease Over Epochs')
    plt.legend()
    plt.grid(True)
    plt.show()


if __name__ == "__main__":
    # 设备选择
    device_config = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # 加载并配置模型
    network = CNN()
    network.to(device=device_config)
    # 指定数据路径并启动训练
    data_directory_path = r"C:\Users\yinjie\Desktop\archive\train"
    train_model(network, device_config, data_directory_path)
